package com.zck.aiagent.advisor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

@Slf4j

public class MyLoggerAdvisor  implements CallAroundAdvisor, StreamAroundAdvisor {

    private int order;


    public MyLoggerAdvisor( int order) {
        this.order = order;
    }

    public MyLoggerAdvisor( ) {
        this.order = 1;
    }


    public String getName() {
        return this.getClass().getSimpleName();
    }

    public int getOrder() {
        return this.order;
    }

    private AdvisedRequest before(AdvisedRequest request) {
        log.info("AI Request: {}", request.userText());
        //在advisor拦截器里面共享消息
//        request = request.updateContext(context -> {
//            context.put("key", "value");
//            return context;
//        });
        return request;
    }

    private void observeAfter(AdvisedResponse advisedResponse) {
        if (advisedResponse.response() != null) {
            log.info("AI Response: {}", advisedResponse.response().getResult().getOutput().getText());
        }
    }

    public String toString() {
        return MyLoggerAdvisor.class.getSimpleName();
    }

    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
        this.observeAfter(advisedResponse);
        return advisedResponse;
    }

    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
        advisedRequest = this.before(advisedRequest);
        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);
        return (new MessageAggregator()).aggregateAdvisedResponse(advisedResponses, this::observeAfter);
    }

    /*对于复杂处理的流,可以使用Reactor的操作符*/
//    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
//        return Mono.just(advisedRequest)
//                .publishOn(Schedulers.boundedElastic())
//                .map(request -> {
//                    // 请求前处理逻辑
//                    return modifyRequest(request);
//                })
//                .flatMapMany(request -> chain.nextAroundStream(request))
//                .map(response -> {
//                    // 响应处理逻辑
//                    return modifyResponse(response);
//                });
//    }

    private AdvisedResponse modifyResponse(AdvisedResponse response) {
        return response;
    }

    private Object modifyRequest(AdvisedRequest request) {
        return  request;
    }

}
